import numpy as np
from matplotlib import pyplot as plt
import matplotlib.cm as cm
import h5py

lx = 0.8 #cm
ly = 0.8 #cm
Nx = 200
Ny = 200

dx = lx/(Nx-1)
dy = ly/(Ny-1)

# opt = 1 : grid density
# opt = 2 : grid velocity x
# opt = 3 : grid velocity y
# opt = 4 : grid pressure
# opt = 5 : grid internal energy
# opt = 6 : grid phi
opt = 4

if opt == 1 :
	# contour fig
	file = h5py.File('../data/rho.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])

	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()



	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()

elif opt == 2 :
	# contour fig
	file = h5py.File('../data/ux.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])

	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()


	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()

elif opt == 3 :
	# contour fig
	file = h5py.File('../data/uy.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])
	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()


	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()

elif opt == 4 :
	# contour fig
	file = h5py.File('../data/p.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])

	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')	
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()



	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()

elif opt == 5 :
	# contour fig
	file = h5py.File('../data/e.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])

	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')	
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()



	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()

elif opt == 6 :
	# contour fig
	file = h5py.File('../data/phi.h5', 'r')
	# print("Keys : %s" % file.keys())
	dsetname = list(file.keys())[0]
	# print(dsetname)

	data = tuple(file[dsetname])
	# print(type(data))
	# print(data[0][0])

	plt.figure(1,figsize=(8,8),dpi=100)
	plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='turbo')
	# plt.imshow(data,interpolation='nearest',origin = 'lower',cmap='nipy_spectral')
	plt.title(dsetname)
	# plt.gca().invert_yaxis() # 有origin = 'lower' 就不用了
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')	
	# plt.xticks(np.arange(0,Nx+1,1),np.arange(-dx,lx+dx,dx))
	# plt.yticks(np.arange(0,Ny+1,1),np.arange(-dy,ly+dy,dy))
	plt.colorbar(drawedges=False)
	plt.show()



	# 3d fig
	x = np.linspace(0, Nx+1, Nx+2)
	y = np.linspace(0, Ny+1, Ny+2)
	z = np.zeros(shape=(Nx+2, Ny+2))
	indx=x.astype('int')
	indy=y.astype('int')
	# print(indx)
	# print(indy)
	for i in indx :
		for j in indy :
			z[i,j] = data[i][j]
	# print(type(z))
	# print(type(x))
	xx, yy=np.meshgrid(x, y)
	x,y,z=xx.ravel(), yy.ravel(), z.ravel()
	plt.figure(2,figsize=(8,8),dpi=100)
	ax = plt.axes(projection='3d')
	ax.bar3d(x,y, np.zeros_like(x),5,5,z,shade=False,color=cm.ScalarMappable().to_rgba(x),cmap='turbo')
	plt.gca().invert_yaxis()
	plt.title(dsetname)
	plt.xlabel('x[cm]')
	plt.ylabel('y[cm]')
	plt.show()




















